
from matrix import *
import sys
import math
import random
import Image
import ImageDraw


CONTRAST_TRH  = 20
DIST_THR      =  7
CROSS_THR     =  4
REMOVE_LENGTH =  5

class Observation:
	def __init__(self, x, p1, p2, lumi):
		self.x = x
		self.position = (p1 + p2) / 2
		self.thickness = abs(p1 - p2)
		self.luminosity = lumi
	def __str__(self):
		return ("(" + str(self.position) + "," + str(self.thickness)
				+ "," + str(self.luminosity) + ")")
	def __repr__(self):
		return str(self)

def process_column(x, pix):
	y = 0
	value = pix[x,y]
	llumi = value
	obs = []
	while y < height:
		while y < height and pix[x, y] >= value:
			value = pix[x, y]
			y = y + 1
		if y >= height:
			break
		bln = y - 1
		blumi = pix[x, bln]
		while y < height and pix[x, y] <= value:
			value = pix[x, y]
			y = y + 1
		if y >= height:
			break
		mln = y - 1
		mlumi = pix[x, mln]
		while y < height and pix[x, y] > value:
			value = pix[x, y]
			y = y + 1
		if y >= height:
			break
		lln = y - 1
		llumi = pix[x, lln]
		y = y + 1
		if (max(blumi, llumi) - mlumi) > CONTRAST_TRH:
			mi1ln = (mln + bln) / 2
			mi2ln = (lln + mln) / 2
			obs.append(Observation(x, mi1ln, mi2ln, mlumi))

	return obs

# x = (position, slope, thickness, luminosity)
# z = (position, thickness, luminosity)
# A = [ 1 1 0 0 ]
#     [ 0 1 0 0 ]
#     [ 0 0 1 0 ]
#     [ 0 0 0 1 ]
# H = [ 1 0 0 0 ]
#     [ 0 0 1 0 ]
#     [ 0 0 0 1 ]
# x_k = A * x_(k-1) + W_(k-1)
# z_k = H * x_(k-1) + V_(k-1)
# R covariance de V => measure covariance
# Q covariance de W => process covariance ~= 10^-5
# x^_k = x~_k + K_k * ( z_k - H * x~_k )

# step 1: (estimate)
# x~_k = A * x^_(k-1)
# P~_k = A P_(k-1)At + Q

# step 2: (update)
#         P~_k * Ht
# K_k = -------------
#       H P~_k Ht + R
# x^_k = x~_k + K_k * ( z_k - H * x~_k )
# P_k = (I - K_k H) P~_k

perr = 10e-7
merr = 10e-7

class Hypothesis:
	def __init__(self, x, position, thickness, luminosity):
		self.points = [ (x, position) ]
		self.x = Matrix([ [position], [0], [thickness], [luminosity] ])
		self.A = Matrix([
			[1, 1, 0, 0],
			[0, 1, 0, 0],
			[0, 0, 1, 0],
			[0, 0, 0, 1] ])
		self.H = Matrix([
			[1, 0, 0, 0],
			[0, 1, 0, 0],
			[0, 0, 1, 0],
			[0, 0, 0, 1] ])
		self.P = Matrix([
			[0, 0, 0, 0],
			[0, 0.0001, 0, 0],
			[0, 0, 0.0001, 0],
			[0, 0, 0, 0.1] ])
		self.Q = Matrix([
			[perr, 0, 0, 0],
			[0, perr, 0, 0],
			[0, 0, perr, 0],
			[0, 0, 0, perr] ])
		self.R = Matrix([
			[merr, 0, 0, 0],
			[0, 0, 0, 0],
			[0, 0, merr, 0],
			[0, 0, 0, merr] ])
		self.removing = 0

	def position(self):
		return self.x[0][0]
	def slope(self):
		return self.x[1][0]
	def thickness(self):
		return self.x[2][0]
	def luminosity(self):
		return self.x[3][0]

	def posCov(self):
		return 1 # self.P[0][0]
	def thickCov(self):
		return 1 # self.P[2][2]
	def lumCov(self):
		return 0.001 # self.P[3][3]

	def distance(self, hyp):
		"mahalanobis distance"
		return math.sqrt(
			((self.position() - hyp.position())*hyp.posCov())**2
			+ ((self.thickness() - hyp.thickness())*hyp.thickCov())**2
			+ ((self.luminosity() - hyp.luminosity())*hyp.lumCov())**2)

	def estimate(self):
		self.xe = self.A * self.x
		self.Pe = self.A * self.P * self.A.transpose() + self.Q
		#print str(self) + " -> xe = " + str(self.xe) + ", Pe=" + str(self.Pe)

	def update(self, ob, xPos):
		self.removing = 0
		z = Matrix([[ob.position], [ob.position-self.position()],
				[ob.thickness], [ob.luminosity]])
		k = self.Pe * self.H.transpose() * self.inv(self.H * self.Pe
				* self.H.transpose() + self.R)
		self.x = self.xe + k * (z - self.H * self.xe)
		self.P = (matrix_I(self.P.size()) - k * self.H) * self.Pe
		#print " -> x = " + str(nx) + ", P=" + str(nP)
		if abs(xPos - self.points[-1][0]) >= 10:
			self.points.append((xPos, self.position()))

	# This is an optimized version of the matrix inversion
	def inv(self, A):
		res = matrix_I(A.size())
		a = A[0][0]
		b = A[0][1]
		c = A[1][0]
		d = A[1][1]
		det = a*d - b*c
		res[0][0] = d / det
		res[0][1] = -b / det
		res[1][0] = -c / det
		res[1][1] = a / det
		res[2][2] = 1.0 / A[2][2]
		res[3][3] = 1.0 / A[3][3]
		return res

	def __repr__(self):
		return str(self)

	def __str__(self):
		return ("(" + str(self.position()) + "," + str(self.thickness())
				+ "," + str(self.luminosity()) + ")")


def distance(ob, hyp):
	"mahalanobis distance"
	return math.sqrt(
		((ob.position - hyp.position())*hyp.posCov())**2
		+ ((ob.thickness - hyp.thickness())*hyp.thickCov())**2
		+ ((ob.luminosity - hyp.luminosity())*hyp.lumCov())**2)


def create_hyp(actives, obs):
	print "create hyp for " + str(obs)
	actives.append(Hypothesis(obs.x, obs.position,
			obs.thickness, obs.luminosity))

def remove_hyp(actives, completed, hyp, x, retry=REMOVE_LENGTH):
	print "remove hyp " + str(hyp)
	if hyp.removing == 0:
		hyp.end = (x-1, hyp.position())
	hyp.removing += 1
	if hyp.removing > retry:
		print "complete hyp " + str(hyp)
		hyp.points.append(hyp.end)
		if hyp in actives:
			actives.remove(hyp)
			completed.append(hyp)
		

def estimate_stage(actives):
	for hyp in actives:
		hyp.estimate()

def matching_stage(actives, completed, obs, x):
	matched_hyp = set([])
	prev_actives = actives[:] # do a copy
	updateMap = {}
	for ob in obs:
		print "matching: " + str(ob) + " at x=" + str(x)
		best = []
		for hyp in prev_actives:
			d = distance(ob, hyp)
			if d < DIST_THR:
				best.append((hyp, d))
		print "best: " + str(best)
		if len(best) == 0:
			create_hyp(actives, ob)
		else:
			last = None
			for (hyp, d) in best:
				if hyp in updateMap.keys():
					(obp, dp) = updateMap[hyp]
					if d < dp:
						print 'other ob:', obp, dp
						updateMap[hyp] = (ob, d)
				else:
					updateMap[hyp] = (ob, d)
				matched_hyp.add(hyp)
				# compute distance between hyps
				# in order to differenciate crossing from joining
				if last and last.distance(hyp) < CROSS_THR:
					print 'cross:', last.distance(hyp)
					h = hyp
					if last.removing > hyp.removing:
						h = last
					remove_hyp(actives, completed, h, x, 3)
				last = hyp

	for hyp in updateMap.keys():
		(ob, d) = updateMap[hyp]
		hyp.update(ob, x)

	for hyp in prev_actives:
		if not hyp in matched_hyp:
			remove_hyp(actives, completed, hyp, x)



#im = Image.open('images/partition.png')
im = Image.open('images/partition_small.png')
#im = Image.open('images/synth.png')
#im = Image.open('images/synth2.png')
#im = Image.open('images/synth3.png')

width, height = im.size
pix = im.load()

actives = []
completed = []
maxObs = 0
nim = im.convert("RGB")
draw = ImageDraw.Draw(nim)
colors = []
for i in range(20):
	colors.append("#%02x%02x%02X" %
		(random.randint(0,255),
		random.randint(0,255),
		random.randint(0,255)))
for x in range(width):
	obs = process_column(x, pix)
	if len(obs) > maxObs:
		maxObs = len(obs)
	for i in range(len(obs)):
		ob = obs[i]
		y1 = ob.position - ob.thickness/2
		y2 = ob.position + ob.thickness/2
		draw.line((x, y1, x, y2), fill=colors[i], width=1)
	print "=== estimate at x=" + str(x)
	estimate_stage(actives)
	print "=== match at x=" + str(x)
	matching_stage(actives, completed, obs, x)

print "max obs: " + str(maxObs)
print "completed:" + str(len(completed))

nim.save("out1.png", "PNG")

nim = im.convert("RGB")
draw = ImageDraw.Draw(nim)
for hyp in completed:
	(x1, y1) = hyp.points[0]
	(x2, y2) = hyp.points[-1]
	if math.sqrt((x1-x2)**2 + (y1-y2)**2) < 10:
		continue
	#print "(" + str(x1) +"," + str(y1) +", " + str(x2) + "," + str(y2) + ")"
	last = hyp.points[0]
	draw.ellipse((last[0]-2, last[1]-2, last[0]+2, last[1]+2), fill="#ff0000")
	for p in hyp.points[1:]:
		draw.line((last[0], last[1], p[0], p[1]), fill="#ff0000", width=1)
		last = p
	draw.ellipse((last[0]-2, last[1]-2, last[0]+2, last[1]+2), fill="#ff0000")

nim.save("out.png", "PNG")
	
